package com.ics.atable.chat.controller.graph;

import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.NodeOutput;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.async.AsyncGenerator;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
import com.alibaba.cloud.ai.graph.state.StateSnapshot;
import com.alibaba.cloud.ai.graph.streaming.StreamingOutput;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.ics.atable.chat.model.dto.*;
import com.ics.atable.chat.model.entity.item.Column;
import com.ics.atable.chat.model.entity.item.Item;
import com.ics.atable.chat.model.entity.item.Table;
import com.ics.atable.chat.model.vo.address.AddressVO;
import com.ics.atable.chat.model.vo.item.ItemVO;
import com.ics.atable.chat.model.vo.node.IntentData;
import com.ics.atable.chat.model.vo.node.NodeResponse;
import com.ics.atable.chat.model.vo.node.RetrievalData;
import com.ics.atable.chat.model.vo.node.StructTokenData;
import com.ics.atable.chat.model.vo.node.stream.ObjectStreamResponse;
import com.ics.atable.chat.model.vo.node.stream.StreamResponse;
import com.ics.atable.chat.model.vo.node.stream.StringStreamResponse;
import com.ics.atable.chat.model.enums.ResponseTypeEnum;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.http.codec.ServerSentEvent;
import reactor.core.publisher.Sinks;

import java.time.LocalDateTime;
import java.time.ZoneId;
import java.util.*;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.StateGraph.START;
import static com.ics.atable.chat.constant.ItemQueryConstant.*;

/**
 * @author haohongbin
 */
@Slf4j
public class GraphProcess {

    private final ExecutorService executor = Executors.newFixedThreadPool(10);

    private final CompiledGraph compiledGraph;

    public GraphProcess(CompiledGraph compiledGraph) {
        this.compiledGraph = compiledGraph;
    }

    public void handleHumanFeedback(ChatRequest chatRequest, Map<String, Object> objectMap,
                                    RunnableConfig runnableConfig, Sinks.Many<ServerSentEvent<StreamResponse>> sink) throws GraphRunnerException {
        objectMap.put("feedback_content", chatRequest.interruptFeedback());
        StateSnapshot stateSnapshot = compiledGraph.getState(runnableConfig);
        OverAllState state = stateSnapshot.state();
        state.withResume();
        state.withHumanFeedback(new OverAllState.HumanFeedback(objectMap, INTENT));
        AsyncGenerator<NodeOutput> resultFuture = compiledGraph.streamFromInitialNode(state, runnableConfig);
        processStream(resultFuture, sink, ResponseTypeEnum.STRING);
    }

    public void processStream(AsyncGenerator<NodeOutput> generator, Sinks.Many<ServerSentEvent<StreamResponse>> sink, ResponseTypeEnum responseTypeEnum) {
        executor.submit(() -> {
            generator.forEachAsync(output -> {
                try {
                    String nodeName = output.node();
                    StreamResponse streamResponse;
                    String content;
                    if (!Objects.equals(nodeName, START) && !Objects.equals(nodeName, END)
                            && !Objects.equals(nodeName, INIT) && !Objects.equals(nodeName, SYNONYMS)) {

                        if (output instanceof StreamingOutput streamingOutput) {
                            // 处理流式输出
                            content = JSON.toJSONString(Map.of(nodeName, streamingOutput.chunk()));
                            streamResponse = new StringStreamResponse(content);
                            log.debug("Streaming output from node {}: {}", nodeName, streamingOutput.chunk());
                        } else {
                            // 处理节点输出
                            streamResponse = handleNodeObjectOutput(output, nodeName, responseTypeEnum);
                        }
                        assert streamResponse != null;
                        sink.tryEmitNext(ServerSentEvent.builder(streamResponse).build());
                    }
                }
                catch (Exception e) {
                    log.error("Error processing output", e);
                    throw new CompletionException(e);
                }
            }).thenAccept(v -> {
                // 正常完成
                sink.tryEmitComplete();
            }).exceptionally(e -> {
                log.error("Error in stream processing", e);
                sink.tryEmitError(e);
                return null;
            });
        });
    }

    /**
     * 处理节点对象输出
     * @param output
     * @param nodeName
     * @param responseTypeEnum
     */
    private StreamResponse handleNodeObjectOutput(NodeOutput output, String nodeName, ResponseTypeEnum responseTypeEnum) {
        if (responseTypeEnum == ResponseTypeEnum.STRING) {
            return handleNodeObjectOutputWithString(output, nodeName);
        } else {
            return handleNodeObjectOutputWithNodeObject(output, nodeName);
        }
    }

    private StreamResponse handleNodeObjectOutputWithNodeObject(NodeOutput output, String nodeName) {
        List<NodeResponse> nodeResponseList = new ArrayList<>();

        switch (nodeName) {
            case INTENT -> {
                NodeResponse nodeResponse = new NodeResponse();
                nodeResponse.setNode("意图识别");

                Optional<Object> intentResultOpt = output.state().value("intent_result");
                IntentData intentData = new IntentData();
                if (intentResultOpt.isPresent()) {
                    Object intentResult = intentResultOpt.get();
                    IntentExtract intentExtract;

                    if (intentResult instanceof IntentExtract) {
                        // 直接转换
                        intentExtract = (IntentExtract) intentResult;
                    } else if (intentResult instanceof Map) {
                        // 从Map转换
                        intentExtract = JSON.parseObject(JSON.toJSONString(intentResult), IntentExtract.class);
                    } else {
                        // 尝试通过JSON转换
                        intentExtract = JSON.parseObject(JSON.toJSONString(intentResult), IntentExtract.class);
                    }

                    intentData.setResult(intentExtract.getOutput());
                } else {
                    throw new IllegalArgumentException("process intent is missing");
                }
                nodeResponse.setData(intentData);
                nodeResponse.setTimestamp(Date.from(LocalDateTime.now().atZone(ZoneId.systemDefault()).toInstant()));
                nodeResponseList.add(nodeResponse);
            }
            case STRUCT_TOKEN -> {
                NodeResponse nodeResponse = new NodeResponse();
                nodeResponse.setNode("结构化提取");

                Optional<Object> structTokenOpt = output.state().value("struct_token");
                StructTokenData structData = new StructTokenData();
                if (structTokenOpt.isPresent()) {
                    Object structToken = structTokenOpt.get();
                    KeyWordExtract keyWordExtract;

                    if (structToken instanceof KeyWordExtract) {
                        // 直接转换
                        keyWordExtract = (KeyWordExtract) structToken;
                    } else if (structToken instanceof Map) {
                        // 从Map转换
                        keyWordExtract = JSON.parseObject(JSON.toJSONString(structToken), KeyWordExtract.class);
                    } else {
                        // 尝试通过JSON转换
                        keyWordExtract = JSON.parseObject(JSON.toJSONString(structToken), KeyWordExtract.class);
                    }
                    structData.setItems(keyWordExtract.getItems());
                    structData.setIndicators(keyWordExtract.getIndicators());
                    structData.setTime(keyWordExtract.getTime());
                    structData.setAddress(keyWordExtract.getAddress());
                } else {
                    throw new IllegalArgumentException("struct_token is missing");
                }
                nodeResponse.setData(structData);
                nodeResponse.setTimestamp(Date.from(LocalDateTime.now().atZone(ZoneId.systemDefault()).toInstant()));
                nodeResponseList.add(nodeResponse);
            }
            case ITEM_RETRIEVAL -> {
                NodeResponse nodeResponse = new NodeResponse();
                nodeResponse.setNode("指标检索");

                RetrievalData retrievalData = new RetrievalData();
                List<ItemVO> itemVOList = new ArrayList<>();
                output.state().value("retrieved_items").ifPresentOrElse(retrievedItems -> {
                    for (Object item : (List<?>) retrievedItems) {
                        ItemVO itemVO = new ItemVO();
                        MatchedItem matchedItem = null;

                        if (item instanceof MatchedItem) {
                            matchedItem = (MatchedItem) item;
                        } else if (item instanceof Map) {
                            matchedItem = JSON.parseObject(JSON.toJSONString(item), MatchedItem.class);
                        } else {
                            matchedItem = JSON.parseObject(JSON.toJSONString(item), MatchedItem.class);
                        }
                        matchedItemConvertToItemVO(itemVO, matchedItem);
                        itemVOList.add(itemVO);
                    }
                    retrievalData.setItemVOList(itemVOList);
                }, () -> {});

                List<AddressVO> addressVOList = new ArrayList<>();
                output.state().value("retrieved_addresses").ifPresentOrElse(retrievedAddresses -> {
                        for (Object address : (List<?>) retrievedAddresses) {
                            AddressVO addressVO = new AddressVO();
                            MatchedAddress matchedAddress = null;

                            if (address instanceof MatchedAddress) {
                                matchedAddress = (MatchedAddress) address;
                            } else if (address instanceof Map) {
                                matchedAddress = JSON.parseObject(JSON.toJSONString(address), MatchedAddress.class);
                            } else {
                                matchedAddress = JSON.parseObject(JSON.toJSONString(address), MatchedAddress.class);
                            }
                            matchedAddressConvertToAddressVO(addressVO, matchedAddress);
                            addressVOList.add(addressVO);
                        }
                        retrievalData.setAddressVOList(addressVOList);
                }, () -> {});

                nodeResponse.setData(retrievalData);
                nodeResponse.setTimestamp(Date.from(LocalDateTime.now().atZone(ZoneId.systemDefault()).toInstant()));
                nodeResponseList.add(nodeResponse);
            }
        }

        return new ObjectStreamResponse(nodeResponseList);
    }


    private StreamResponse handleNodeObjectOutputWithString(NodeOutput output, String nodeName) {
        JSONObject nodeOutput = new JSONObject();
        switch (nodeName) {
            case INTENT -> {
                nodeOutput.put("node", nodeName);

                Optional<Object> intentResultOpt = output.state().value("intent_result");
                if (intentResultOpt.isPresent()) {
                    Object intentResult = intentResultOpt.get();
                    IntentExtract intentExtract;

                    if (intentResult instanceof IntentExtract) {
                        // 直接转换
                        intentExtract = (IntentExtract) intentResult;
                    } else if (intentResult instanceof Map) {
                        // 从Map转换
                        intentExtract = JSON.parseObject(JSON.toJSONString(intentResult), IntentExtract.class);
                    } else {
                        // 尝试通过JSON转换
                        intentExtract = JSON.parseObject(JSON.toJSONString(intentResult), IntentExtract.class);
                    }

                    nodeOutput.put("意图识别结果：", intentExtract.getOutput());
                } else {
                    throw new IllegalArgumentException("process intent is missing");
                }
            }
            case STRUCT_TOKEN -> {
                nodeOutput.put("node", nodeName);

                Optional<Object> structTokenOpt = output.state().value("struct_token");
                if (structTokenOpt.isPresent()) {
                    Object structToken = structTokenOpt.get();
                    KeyWordExtract keyWordExtract;

                    if (structToken instanceof KeyWordExtract) {
                        // 直接转换
                        keyWordExtract = (KeyWordExtract) structToken;
                    } else if (structToken instanceof Map) {
                        // 从Map转换
                        keyWordExtract = JSON.parseObject(JSON.toJSONString(structToken), KeyWordExtract.class);
                    } else {
                        // 尝试通过JSON转换
                        keyWordExtract = JSON.parseObject(JSON.toJSONString(structToken), KeyWordExtract.class);
                    }

                    // 创建一个 JSONObject 来表示结构化提取结果
                    JSONObject structResult = new JSONObject();
                    structResult.put("统计项", keyWordExtract.getItems());
                    structResult.put("统计指标", keyWordExtract.getIndicators());
                    structResult.put("时间", keyWordExtract.getTime());
                    structResult.put("地址", keyWordExtract.getAddress());

                    nodeOutput.put("结构化提取结果", structResult);
                } else {
                    throw new IllegalArgumentException("struct_token is missing");
                }
            }
            case ITEM_RETRIEVAL -> {
                nodeOutput.put("node", nodeName);

                output.state().value("retrieved_items").ifPresentOrElse(retrievedItems -> {
                    if (retrievedItems instanceof List) {
                        StringBuilder sb = new StringBuilder();
                        for (Object item : (List<?>) retrievedItems) {
                            MatchedItem matchedItem = null;

                            if (item instanceof MatchedItem) {
                                matchedItem = (MatchedItem) item;
                            } else if (item instanceof Map) {
                                matchedItem = JSON.parseObject(JSON.toJSONString(item), MatchedItem.class);
                            } else {
                                matchedItem = JSON.parseObject(JSON.toJSONString(item), MatchedItem.class);
                            }

                            if (matchedItem != null) {
                                sb.append(matchedItem.getItemFullName());
                                Set<String> itemSet = new LinkedHashSet<>(); // 使用LinkedHashSet保持插入顺序并去重
                                List<Table> tableList = matchedItem.getTableList();
                                for (Table table : tableList) {
                                    table.getColumnList().stream()
                                            .map(Column::getColumnName)
                                            .forEach(itemSet::add);
                                }
                                List<String> itemList = new ArrayList<>(itemSet);
                                if (CollectionUtils.isNotEmpty(itemList)) {
                                    sb.append("[");
                                    for (int i = 0; i < itemList.size(); i++) {
                                        if (i > 0) sb.append(", ");
                                        sb.append(itemList.get(i));
                                    }
                                    sb.append("]");
                                }
                                sb.append("、");
                            }
                        }
                        if (!sb.isEmpty()) {
                            sb.deleteCharAt(sb.length() - 1);
                        }
                        nodeOutput.put("统计项[统计指标]", sb.toString());
                    } else {
                        nodeOutput.put("统计项[统计指标]", retrievedItems);
                    }
                }, () -> {
                    nodeOutput.put("统计项[统计指标]", "未找到检索结果");
                });

                output.state().value("retrieved_addresses").ifPresentOrElse(retrievedAddresses -> {
                    if (retrievedAddresses instanceof List) {
                        StringBuilder sb = new StringBuilder();
                        for (Object address : (List<?>) retrievedAddresses) {
                            MatchedAddress matchedAddress = null;

                            if (address instanceof MatchedAddress) {
                                matchedAddress = (MatchedAddress) address;
                            } else if (address instanceof Map) {
                                matchedAddress = JSON.parseObject(JSON.toJSONString(address), MatchedAddress.class);
                            } else {
                                matchedAddress = JSON.parseObject(JSON.toJSONString(address), MatchedAddress.class);
                            }
                            if (matchedAddress != null) {
                                List<String> addressParts = Arrays.asList(
                                        matchedAddress.getProvince(),
                                        matchedAddress.getCity(),
                                        matchedAddress.getCounty(),
                                        matchedAddress.getTown(),
                                        matchedAddress.getVillage()
                                );
                                addressParts.stream()
                                        .filter(Objects::nonNull)
                                        .filter(part -> !part.isEmpty())
                                        .forEach(sb::append);
                            }
                            sb.append("、");
                        }
                        if (!sb.isEmpty()) {
                            sb.deleteCharAt(sb.length() - 1);
                        }
                        nodeOutput.put("地址", sb.toString());
                    } else {
                        nodeOutput.put("地址", retrievedAddresses);
                    }
                }, () -> {
                    nodeOutput.put("地址", "未匹配到地址");
                });
            }
        }

        return new StringStreamResponse(JSON.toJSONString(nodeOutput));
    }

    /**
     * 将 MatchedAddress 转换为 AddressVO
     * @param addressVO addressVO
     * @param matchedAddress matchedAddress
     */
    private void matchedAddressConvertToAddressVO(AddressVO addressVO, MatchedAddress matchedAddress) {
        if (matchedAddress != null) {
            addressVO.setOrgId(matchedAddress.getOrgId());
            addressVO.setVillageName(matchedAddress.getVillage());
            addressVO.setTownName(matchedAddress.getTown());
            addressVO.setCountyName(matchedAddress.getCounty());
            addressVO.setCityName(matchedAddress.getCity());
            addressVO.setProvinceName(matchedAddress.getProvince());
            addressVO.setSimilarity(matchedAddress.getScore());
        }
    }

    /**
     * 将 MatchedItem 转换为 ItemVO
     * @param itemVO itemVO
     * @param matchedItem matchedItem
     */
    private void matchedItemConvertToItemVO(ItemVO itemVO, MatchedItem matchedItem) {
        if (matchedItem != null) {
            itemVO.setItem(new Item());
            itemVO.getItem().setItemUuid(matchedItem.getItemUuid());
            itemVO.getItem().setItemName(matchedItem.getItemName());
            itemVO.getItem().setItemFullName(matchedItem.getItemFullName());
            itemVO.getItem().setItemFatherName(matchedItem.getItemFatherName());
            itemVO.getItem().setItemGrandfatherName(matchedItem.getItemGrandfatherName());
            itemVO.getItem().setMetadata(matchedItem.getMetadata());
            itemVO.getItem().setTableList(matchedItem.getTableList());
            itemVO.setSimilarity(matchedItem.getScore());
        }
    }

}
